from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from PIL import Image
import cv2
import numpy as np


if __name__ == "__main__":

    file = "data/train/ants_image/6240329_72c01e663e.jpg"

    writer = SummaryWriter("demo03/logs")

    # 使用 PIL 读取图片文件
    image_origin = Image.open(file)
   # writer.add_image("PIL-Image-Origin", image_origin, 1, dataformats="HWC")

    print("image 类型：", type(image_origin), image_origin)
    writer.add_image("PIL-Image-Origin", np.array(image_origin), 1, dataformats="HWC")

    to_tensor = transforms.ToTensor()

    # 将 PIL 的图片文件转换为 tensor 图片文件
    image_tensor = to_tensor(image_origin)

    print("tensor 类型：", type(image_tensor), image_tensor)
    writer.add_image("PIL-Image-Tensor", image_tensor)

    # 使用 OpenCV 读取图片文件
    image_origin = cv2.imread(file)

    print("image 类型：", type(image_origin), image_origin)
    writer.add_image("CV2-Image-Origin", image_origin, 1, dataformats="HWC")

    # 将 NumPy 的 ndarray 类型的图片文件转换为 tensor 图片文件
    image_tensor = to_tensor(image_origin)

    print("tensor 类型：", type(image_tensor), image_tensor)
    writer.add_image("CV2-Image-Tensor", image_tensor)

    writer.close()


